import os, sys
import math
from functools import reduce
from collections import OrderedDict
from typing import *

import mindspore
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops

import numpy as np
from mindspore import Tensor
from mindspore.ops import operations as P
from mindspore.common import initializer as init

from wpgpmfm.unet3d.networks.generalcells import ForwardConv, MaxPool3d


class Unet3d(nn.Cell):
    """
    the unet3d architecture
    """
    __constants__ = ['features']

    def __init__(self, in_channels, num_classes=2):
        super(Unet3d, self).__init__()
        self.fl1 = ForwardConv(in_channels, 8, 3, 3)
        self.down1 = MaxPool3d(2, 2)

        self.fl2 = ForwardConv(8, 16, 3, 3)
        self.down2 = MaxPool3d(2, 2)

        self.fl3 = ForwardConv(16, 32, 3, 3)

        self.up2 = nn.Conv3dTranspose(in_channels=32, out_channels=16, kernel_size=2, stride=2, pad_mode='same')
        self.fr2 = ForwardConv(16, 16, 3, 3)

        self.up1 = nn.Conv3dTranspose(in_channels=16, out_channels=8, kernel_size=2, stride=2, pad_mode='same')
        self.fr1 = ForwardConv(8, 8, 3, 3)

        self.out = nn.Conv3d(in_channels=8, out_channels=num_classes, kernel_size=1, stride=1, pad_mode='same')

    def construct(self, x):
        x1 = self.fl1(x)
        x2 = self.fl2(self.down1(x1))
        x3 = self.fl3(self.down2(x2))

        x2 = self.fr2(self.up2(x3) + x2)
        x1 = self.fr1(self.up1(x2) + x1)

        x = self.out(x1)
        return x


if __name__ == '__main__':
    #  sorting files:

    ROOT = '/share_data/liupan/hematoma/npyfiles/24h/train'
    x = Tensor(np.ones([2, 1, 32, 32, 32]), mindspore.float32)
    unet3d = Unet3d(1, 2)
    output = unet3d(x)
    print(output.shape)

    # conv3d = nn.Conv3d(in_channels=1, out_channels=2, kernel_size=3, stride=1, pad_mode='same')
    # output = conv3d(x)
    # print(output.shape)
